{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3c4c7d5f-51fb-4879-8fd3-d304165ffd38",
   "metadata": {},
   "source": [
    "# Evaluate bo767 retrieval recall accuracy with NV-Ingest and Milvus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a453802-83f4-4fa5-95f2-b663dfeec59b",
   "metadata": {},
   "source": [
    "In this notebook, we'll use NV-ingest and LlamaIndex to get the end-to-end recall accuracy of a retrieval pipeline made up of NV-Ingest's extraction and embedding tasks and a Milvus vector database (VDB)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "1c174e25-ffdf-4764-bad5-e3be8cb00943",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import MilvusClient\n",
    "\n",
    "milvus_client = MilvusClient(\"http://localhost:19530\")\n",
    "milvus_client.drop_collection(collection_name='text')\n",
    "milvus_client.drop_collection(collection_name='tables')\n",
    "milvus_client.drop_collection(collection_name='charts')\n",
    "milvus_client.drop_collection(collection_name='multimodal')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00aeacea-2eb7-45a8-8e62-edf52fdc9d9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nv_ingest_client.client import Ingestor\n",
    "\n",
    "ingestor = (\n",
    "    Ingestor(message_client_hostname=\"localhost\")\n",
    "    .files(\"../data/bo767/*.pdf\")\n",
    "    .extract(\n",
    "        extract_text=True,\n",
    "        extract_tables=True,\n",
    "        extract_charts=True,\n",
    "        extract_images=False,\n",
    "        text_depth=\"page\",\n",
    "    ).embed()\n",
    ")\n",
    "\n",
    "results = ingestor.ingest()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f1dba08-c468-425f-9eb3-48fe568b67c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: save results\n",
    "import pickle\n",
    "\n",
    "filehandler = open('bo767_results.obj', 'wb')\n",
    "pickle.dump(results, filehandler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f4cd3db7-c8a4-478e-9b48-3c8fffe4d32c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: load results\n",
    "import pickle\n",
    "\n",
    "filehandler = open('bo767_results.obj', 'rb')\n",
    "results = pickle.load(filehandler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "e25582b6-005b-47d2-8b47-b0823422bda9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "767"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "02c25cb3-f912-48fd-8e23-ea5c7288d83a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nv_ingest_client.util.milvus import write_to_nvingest_collection, create_nvingest_collection, nvingest_retrieval\n",
    "\n",
    "sparse = False\n",
    "milvus_hostname = \"localhost\"\n",
    "create_nvingest_collection(\"text\", f\"http://{milvus_hostname}:19530\", sparse=sparse, gpu_search=True)\n",
    "create_nvingest_collection(\"tables\", f\"http://{milvus_hostname}:19530\", sparse=sparse, gpu_search=True)\n",
    "create_nvingest_collection(\"charts\", f\"http://{milvus_hostname}:19530\", sparse=sparse, gpu_search=True)\n",
    "create_nvingest_collection(\"multimodal\", f\"http://{milvus_hostname}:19530\", sparse=sparse, gpu_search=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "0331fe6c-628e-4066-9277-643eb12f46ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_results = [[element for element in results if element['document_type'] == 'text'] for results in results]\n",
    "table_results = [[element for element in results if element['metadata']['content_metadata']['subtype'] == 'table'] for results in results]\n",
    "chart_results = [[element for element in results if element['metadata']['content_metadata']['subtype'] == 'chart'] for results in results]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "21a86f1c-10de-4741-aa9b-d56ea1de84f9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote data to: [['0ca30e66-ca9e-4875-bab6-66535bee39ea/1.parquet']]\n",
      "Start time: 2025-01-24 06:17:49\n",
      "Imported row count: 45816\n",
      "Bulk text upload took 54.15791320800781 s\n",
      "Wrote data to: [['7f738ffc-7fd5-4ce3-86dd-b9ab9c515ddc/1.parquet']]\n",
      "Start time: 2025-01-24 06:19:11\n",
      "Imported row count: 27193\n",
      "Bulk tables upload took 38.110397815704346 s\n",
      "Wrote data to: [['e5460db7-5d86-47bc-981b-46c810cc9d45/1.parquet']]\n",
      "Start time: 2025-01-24 06:20:12\n",
      "Imported row count: 5667\n",
      "Bulk charts upload took 17.05130910873413 s\n",
      "Wrote data to: [['c61b35ef-4224-467d-ba98-fd5e5659b6b1/1.parquet']]\n",
      "Start time: 2025-01-24 06:21:10\n",
      "Imported row count: 78676\n",
      "Bulk multimodal upload took 87.26382780075073 s\n"
     ]
    }
   ],
   "source": [
    "write_to_nvingest_collection(text_results, \"text\", sparse=sparse, milvus_uri=f\"http://{milvus_hostname}:19530\", minio_endpoint=\"localhost:9000\")\n",
    "write_to_nvingest_collection(table_results, \"tables\", sparse=sparse, milvus_uri=f\"http://{milvus_hostname}:19530\", minio_endpoint=\"localhost:9000\")\n",
    "write_to_nvingest_collection(chart_results, \"charts\", sparse=sparse, milvus_uri=f\"http://{milvus_hostname}:19530\", minio_endpoint=\"localhost:9000\")\n",
    "write_to_nvingest_collection(results, \"multimodal\", sparse=sparse, milvus_uri=f\"http://{milvus_hostname}:19530\", minio_endpoint=\"localhost:9000\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "2667d436-0e13-4539-b679-2c922b6069a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "def get_recall_scores(query_df, collection_name):\n",
    "    hits = defaultdict(list)\n",
    "    all_answers = nvingest_retrieval(\n",
    "        df_query[\"query\"].to_list(),\n",
    "        collection_name,\n",
    "        hybrid=sparse,\n",
    "        embedding_endpoint=\"http://localhost:8012/v1\",\n",
    "        model_name=\"nvidia/llama-3.2-nv-embedqa-1b-v2\",\n",
    "        top_k=10,\n",
    "        gpu_search=True,\n",
    "    )\n",
    "\n",
    "    for i in range(len(df_query)):\n",
    "        expected_pdf_page = query_df['pdf_page'][i]\n",
    "        retrieved_answers = all_answers[i]\n",
    "        retrieved_pdfs = [os.path.basename(result['entity']['source']['source_id']).split('.')[0] for result in retrieved_answers]\n",
    "        retrieved_pages = [str(result['entity']['content_metadata']['page_number']) for result in retrieved_answers]\n",
    "        retrieved_pdf_pages = [f\"{pdf}_{page}\" for pdf, page in zip(retrieved_pdfs, retrieved_pages)]    \n",
    "\n",
    "        for k in [1, 3, 5, 10]:\n",
    "            hits[k].append(expected_pdf_page in retrieved_pdf_pages[:k])\n",
    "    \n",
    "    for k in hits:\n",
    "        print(f'  - Recall @{k}: {np.mean(hits[k]) :.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73aa9968-936f-43dc-80e4-8e9b8c649292",
   "metadata": {},
   "source": [
    "## Text Recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "b329ee8b-5e19-4492-8742-18d0606f539a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pdf</th>\n",
       "      <th>query</th>\n",
       "      <th>answer</th>\n",
       "      <th>gt_page</th>\n",
       "      <th>pdf_page</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1102434</td>\n",
       "      <td>How much was the ARtillery Intelligence projec...</td>\n",
       "      <td>$4.2 billion</td>\n",
       "      <td>19</td>\n",
       "      <td>1102434_19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1102434</td>\n",
       "      <td>How much revenue of AR advertising is expected...</td>\n",
       "      <td>$8.8 billion</td>\n",
       "      <td>3</td>\n",
       "      <td>1102434_3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1096078</td>\n",
       "      <td>What types of statistics were utilized by Rein...</td>\n",
       "      <td>descriptive statistics</td>\n",
       "      <td>3</td>\n",
       "      <td>1096078_3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1054125</td>\n",
       "      <td>What was the maximum amount requested for cond...</td>\n",
       "      <td>$35,000.00</td>\n",
       "      <td>1</td>\n",
       "      <td>1054125_1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1246906</td>\n",
       "      <td>What is the median household income for the Ci...</td>\n",
       "      <td>$53,278</td>\n",
       "      <td>7</td>\n",
       "      <td>1246906_7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>483</th>\n",
       "      <td>2089825</td>\n",
       "      <td>Under the Climate Action and Low Carbon Develo...</td>\n",
       "      <td>Denis Naughten TD</td>\n",
       "      <td>0</td>\n",
       "      <td>2089825_0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>484</th>\n",
       "      <td>2089825</td>\n",
       "      <td>How many organizations make up Stop Climate Ch...</td>\n",
       "      <td>30</td>\n",
       "      <td>5</td>\n",
       "      <td>2089825_5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>485</th>\n",
       "      <td>2098077</td>\n",
       "      <td>What is the maximum length of Sai Yok bent-­to...</td>\n",
       "      <td>2.4 inches</td>\n",
       "      <td>1</td>\n",
       "      <td>2098077_1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>486</th>\n",
       "      <td>2098077</td>\n",
       "      <td>What characteristic sets the Sai Yok Bent-toed...</td>\n",
       "      <td>enlarged thigh scales</td>\n",
       "      <td>1</td>\n",
       "      <td>2098077_1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>487</th>\n",
       "      <td>2098077</td>\n",
       "      <td>What is the name of the new gecko species that...</td>\n",
       "      <td>Sai Yok bent­toed gecko ( Cyrtodactylus saiyok )</td>\n",
       "      <td>0</td>\n",
       "      <td>2098077_0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>488 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         pdf                                              query  \\\n",
       "0    1102434  How much was the ARtillery Intelligence projec...   \n",
       "1    1102434  How much revenue of AR advertising is expected...   \n",
       "2    1096078  What types of statistics were utilized by Rein...   \n",
       "3    1054125  What was the maximum amount requested for cond...   \n",
       "4    1246906  What is the median household income for the Ci...   \n",
       "..       ...                                                ...   \n",
       "483  2089825  Under the Climate Action and Low Carbon Develo...   \n",
       "484  2089825  How many organizations make up Stop Climate Ch...   \n",
       "485  2098077  What is the maximum length of Sai Yok bent-­to...   \n",
       "486  2098077  What characteristic sets the Sai Yok Bent-toed...   \n",
       "487  2098077  What is the name of the new gecko species that...   \n",
       "\n",
       "                                               answer  gt_page    pdf_page  \n",
       "0                                        $4.2 billion       19  1102434_19  \n",
       "1                                        $8.8 billion        3   1102434_3  \n",
       "2                              descriptive statistics        3   1096078_3  \n",
       "3                                          $35,000.00        1   1054125_1  \n",
       "4                                             $53,278        7   1246906_7  \n",
       "..                                                ...      ...         ...  \n",
       "483                                 Denis Naughten TD        0   2089825_0  \n",
       "484                                                30        5   2089825_5  \n",
       "485                                        2.4 inches        1   2098077_1  \n",
       "486                             enlarged thigh scales        1   2098077_1  \n",
       "487  Sai Yok bent­toed gecko ( Cyrtodactylus saiyok )        0   2098077_0  \n",
       "\n",
       "[488 rows x 5 columns]"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "df_query = pd.read_csv('../data/text_query_answer_gt_page.csv')\n",
    "df_query.pdf = df_query.pdf.apply(lambda x: x.replace('.pdf',''))\n",
    "df_query['pdf_page'] = df_query.apply(lambda x: f\"{x.pdf}_{x.gt_page}\", axis=1) \n",
    "df_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "00eaf5d1-6fce-4fc9-99c4-935b676d022e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  - Recall @1: 0.627\n",
      "  - Recall @3: 0.826\n",
      "  - Recall @5: 0.877\n",
      "  - Recall @10: 0.914\n"
     ]
    }
   ],
   "source": [
    "get_recall_scores(df_query, \"text\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ee8b8e-b4e0-4367-b1da-7c21f59e8b6d",
   "metadata": {},
   "source": [
    "## Table recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "01483767-680d-4ce9-94b4-5c3c16ebd091",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>pdf</th>\n",
       "      <th>page</th>\n",
       "      <th>table</th>\n",
       "      <th>pdf_page</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>How much did Pendleton County spend out of the...</td>\n",
       "      <td>1003421</td>\n",
       "      <td>2</td>\n",
       "      <td>1003421_2_0</td>\n",
       "      <td>1003421_2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>How many units are occupied by single families...</td>\n",
       "      <td>1008059</td>\n",
       "      <td>6</td>\n",
       "      <td>1008059_6_1</td>\n",
       "      <td>1008059_6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>In the Klamath county, what is the total valua...</td>\n",
       "      <td>1008059</td>\n",
       "      <td>6</td>\n",
       "      <td>1008059_6_1</td>\n",
       "      <td>1008059_6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>How much did Nalco pay GRIDCO for electricity ...</td>\n",
       "      <td>1011810</td>\n",
       "      <td>21</td>\n",
       "      <td>1011810_21_0</td>\n",
       "      <td>1011810_21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>How much coal is used at Alumina refinery of N...</td>\n",
       "      <td>1011810</td>\n",
       "      <td>21</td>\n",
       "      <td>1011810_21_2</td>\n",
       "      <td>1011810_21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>230</th>\n",
       "      <td>How much is the rental income from water plant...</td>\n",
       "      <td>2407280</td>\n",
       "      <td>30</td>\n",
       "      <td>2407280_30_0</td>\n",
       "      <td>2407280_30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>231</th>\n",
       "      <td>In 2020 how much were the supplemental taxes f...</td>\n",
       "      <td>2415001</td>\n",
       "      <td>65</td>\n",
       "      <td>not detected</td>\n",
       "      <td>2415001_65</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>232</th>\n",
       "      <td>As of 2020, what is the total of collections a...</td>\n",
       "      <td>2415001</td>\n",
       "      <td>65</td>\n",
       "      <td>not detected</td>\n",
       "      <td>2415001_65</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>233</th>\n",
       "      <td>What was the net gain from the operations of t...</td>\n",
       "      <td>2416020</td>\n",
       "      <td>84</td>\n",
       "      <td>2416020_84_0</td>\n",
       "      <td>2416020_84</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>234</th>\n",
       "      <td>What was the highest operating expense in Nove...</td>\n",
       "      <td>2416020</td>\n",
       "      <td>84</td>\n",
       "      <td>2416020_84_0</td>\n",
       "      <td>2416020_84</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>235 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 query      pdf  page  \\\n",
       "0    How much did Pendleton County spend out of the...  1003421     2   \n",
       "1    How many units are occupied by single families...  1008059     6   \n",
       "2    In the Klamath county, what is the total valua...  1008059     6   \n",
       "3    How much did Nalco pay GRIDCO for electricity ...  1011810    21   \n",
       "4    How much coal is used at Alumina refinery of N...  1011810    21   \n",
       "..                                                 ...      ...   ...   \n",
       "230  How much is the rental income from water plant...  2407280    30   \n",
       "231  In 2020 how much were the supplemental taxes f...  2415001    65   \n",
       "232  As of 2020, what is the total of collections a...  2415001    65   \n",
       "233  What was the net gain from the operations of t...  2416020    84   \n",
       "234  What was the highest operating expense in Nove...  2416020    84   \n",
       "\n",
       "            table    pdf_page  \n",
       "0     1003421_2_0   1003421_2  \n",
       "1     1008059_6_1   1008059_6  \n",
       "2     1008059_6_1   1008059_6  \n",
       "3    1011810_21_0  1011810_21  \n",
       "4    1011810_21_2  1011810_21  \n",
       "..            ...         ...  \n",
       "230  2407280_30_0  2407280_30  \n",
       "231  not detected  2415001_65  \n",
       "232  not detected  2415001_65  \n",
       "233  2416020_84_0  2416020_84  \n",
       "234  2416020_84_0  2416020_84  \n",
       "\n",
       "[235 rows x 5 columns]"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_query = pd.read_csv('../data/table_queries_cleaned_235.csv')[['query','pdf','page','table']]\n",
    "df_query['pdf_page'] = df_query.apply(lambda x: f\"{x.pdf}_{x.page}\", axis=1)\n",
    "df_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "4d9e3602-f39d-4256-98d7-7772da1f00ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  - Recall @1: 0.502\n",
      "  - Recall @3: 0.732\n",
      "  - Recall @5: 0.787\n",
      "  - Recall @10: 0.855\n"
     ]
    }
   ],
   "source": [
    "get_recall_scores(df_query, \"tables\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f514724a-70f4-4d06-a929-0d2bb848902c",
   "metadata": {},
   "source": [
    "## Chart Recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "ddf8ba20-2de5-4b92-a87f-d78c341c07c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>pdf</th>\n",
       "      <th>page</th>\n",
       "      <th>pdf_page</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What are the top three consumer complaint cate...</td>\n",
       "      <td>1009210</td>\n",
       "      <td>11</td>\n",
       "      <td>1009210_11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Which 3 categories did extremely well in terms...</td>\n",
       "      <td>1009210</td>\n",
       "      <td>11</td>\n",
       "      <td>1009210_11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>What's the longest recent US recession?</td>\n",
       "      <td>1010876</td>\n",
       "      <td>0</td>\n",
       "      <td>1010876_0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Is the 12-Month default rate usually higher th...</td>\n",
       "      <td>1010876</td>\n",
       "      <td>0</td>\n",
       "      <td>1010876_0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Which allegation is submitted highest to RTAs ...</td>\n",
       "      <td>1014669</td>\n",
       "      <td>0</td>\n",
       "      <td>1014669_0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>263</th>\n",
       "      <td>After the 2008 recession, what percentage of p...</td>\n",
       "      <td>2384395</td>\n",
       "      <td>6</td>\n",
       "      <td>2384395_6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>264</th>\n",
       "      <td>what were the top 3 major religious groups in ...</td>\n",
       "      <td>2392676</td>\n",
       "      <td>5</td>\n",
       "      <td>2392676_5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>265</th>\n",
       "      <td>What percentage of people in the world identif...</td>\n",
       "      <td>2392676</td>\n",
       "      <td>5</td>\n",
       "      <td>2392676_5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>266</th>\n",
       "      <td>Between 2003 and 2019, has the household mortg...</td>\n",
       "      <td>2410699</td>\n",
       "      <td>189</td>\n",
       "      <td>2410699_189</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>267</th>\n",
       "      <td>When did the total household mortgage debt in ...</td>\n",
       "      <td>2410699</td>\n",
       "      <td>189</td>\n",
       "      <td>2410699_189</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>268 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 query      pdf  page  \\\n",
       "0    What are the top three consumer complaint cate...  1009210    11   \n",
       "1    Which 3 categories did extremely well in terms...  1009210    11   \n",
       "2              What's the longest recent US recession?  1010876     0   \n",
       "3    Is the 12-Month default rate usually higher th...  1010876     0   \n",
       "4    Which allegation is submitted highest to RTAs ...  1014669     0   \n",
       "..                                                 ...      ...   ...   \n",
       "263  After the 2008 recession, what percentage of p...  2384395     6   \n",
       "264  what were the top 3 major religious groups in ...  2392676     5   \n",
       "265  What percentage of people in the world identif...  2392676     5   \n",
       "266  Between 2003 and 2019, has the household mortg...  2410699   189   \n",
       "267  When did the total household mortgage debt in ...  2410699   189   \n",
       "\n",
       "        pdf_page  \n",
       "0     1009210_11  \n",
       "1     1009210_11  \n",
       "2      1010876_0  \n",
       "3      1010876_0  \n",
       "4      1014669_0  \n",
       "..           ...  \n",
       "263    2384395_6  \n",
       "264    2392676_5  \n",
       "265    2392676_5  \n",
       "266  2410699_189  \n",
       "267  2410699_189  \n",
       "\n",
       "[268 rows x 4 columns]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_query = pd.read_csv('../data/charts_with_page_num_fixed.csv')[['query','pdf','page']]\n",
    "df_query['page'] = df_query['page']-1 # page -1 because the page number starts with 1 in that csv\n",
    "df_query['pdf_page'] = df_query.apply(lambda x: f\"{x.pdf}_{x.page}\", axis=1) \n",
    "df_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "bae55dda-e491-4de1-8b98-da0bf187e0e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  - Recall @1: 0.612\n",
      "  - Recall @3: 0.743\n",
      "  - Recall @5: 0.795\n",
      "  - Recall @10: 0.817\n"
     ]
    }
   ],
   "source": [
    "get_recall_scores(df_query, \"charts\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3a0a35d-dd4d-4d47-9f39-b3823b87d44f",
   "metadata": {},
   "source": [
    "## Multimodal Recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "93d5aa3c-6a3e-4c57-b22e-396d0bbed50d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>pdf</th>\n",
       "      <th>page</th>\n",
       "      <th>modality</th>\n",
       "      <th>pdf_page</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>How much was the ARtillery Intelligence projec...</td>\n",
       "      <td>1102434</td>\n",
       "      <td>19</td>\n",
       "      <td>text</td>\n",
       "      <td>1102434_19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>How much revenue of AR advertising is expected...</td>\n",
       "      <td>1102434</td>\n",
       "      <td>3</td>\n",
       "      <td>text</td>\n",
       "      <td>1102434_3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>What types of statistics were utilized by Rein...</td>\n",
       "      <td>1096078</td>\n",
       "      <td>3</td>\n",
       "      <td>text</td>\n",
       "      <td>1096078_3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What was the maximum amount requested for cond...</td>\n",
       "      <td>1054125</td>\n",
       "      <td>1</td>\n",
       "      <td>text</td>\n",
       "      <td>1054125_1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What is the median household income for the Ci...</td>\n",
       "      <td>1246906</td>\n",
       "      <td>7</td>\n",
       "      <td>text</td>\n",
       "      <td>1246906_7</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>986</th>\n",
       "      <td>After the 2008 recession, what percentage of p...</td>\n",
       "      <td>2384395</td>\n",
       "      <td>6</td>\n",
       "      <td>chart</td>\n",
       "      <td>2384395_6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>987</th>\n",
       "      <td>what were the top 3 major religious groups in ...</td>\n",
       "      <td>2392676</td>\n",
       "      <td>5</td>\n",
       "      <td>chart</td>\n",
       "      <td>2392676_5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>988</th>\n",
       "      <td>What percentage of people in the world identif...</td>\n",
       "      <td>2392676</td>\n",
       "      <td>5</td>\n",
       "      <td>chart</td>\n",
       "      <td>2392676_5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>989</th>\n",
       "      <td>Between 2003 and 2019, has the household mortg...</td>\n",
       "      <td>2410699</td>\n",
       "      <td>189</td>\n",
       "      <td>chart</td>\n",
       "      <td>2410699_189</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>990</th>\n",
       "      <td>When did the total household mortgage debt in ...</td>\n",
       "      <td>2410699</td>\n",
       "      <td>189</td>\n",
       "      <td>chart</td>\n",
       "      <td>2410699_189</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>991 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 query      pdf  page  \\\n",
       "0    How much was the ARtillery Intelligence projec...  1102434    19   \n",
       "1    How much revenue of AR advertising is expected...  1102434     3   \n",
       "2    What types of statistics were utilized by Rein...  1096078     3   \n",
       "3    What was the maximum amount requested for cond...  1054125     1   \n",
       "4    What is the median household income for the Ci...  1246906     7   \n",
       "..                                                 ...      ...   ...   \n",
       "986  After the 2008 recession, what percentage of p...  2384395     6   \n",
       "987  what were the top 3 major religious groups in ...  2392676     5   \n",
       "988  What percentage of people in the world identif...  2392676     5   \n",
       "989  Between 2003 and 2019, has the household mortg...  2410699   189   \n",
       "990  When did the total household mortgage debt in ...  2410699   189   \n",
       "\n",
       "    modality     pdf_page  \n",
       "0       text   1102434_19  \n",
       "1       text    1102434_3  \n",
       "2       text    1096078_3  \n",
       "3       text    1054125_1  \n",
       "4       text    1246906_7  \n",
       "..       ...          ...  \n",
       "986    chart    2384395_6  \n",
       "987    chart    2392676_5  \n",
       "988    chart    2392676_5  \n",
       "989    chart  2410699_189  \n",
       "990    chart  2410699_189  \n",
       "\n",
       "[991 rows x 5 columns]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_query = pd.read_csv('../data/text_query_answer_gt_page.csv').rename(columns={'gt_page':'page'})[['query','pdf','page']]\n",
    "df_query.pdf = df_query.pdf.apply(lambda x: x.replace('.pdf',''))\n",
    "df_query['modality'] = 'text'\n",
    "\n",
    "df_query2 = pd.read_csv('../data/table_queries_cleaned_235.csv')[['query','pdf','page']]\n",
    "df_query2['modality'] = 'table'\n",
    "\n",
    "df_query3 = pd.read_csv('../data/charts_with_page_num_fixed.csv')[['query','pdf','page']]\n",
    "df_query3['page'] = df_query3['page']-1 # page -1 because the page number starts with 1 in that csv\n",
    "df_query3['modality'] = 'chart'\n",
    "\n",
    "df_query = pd.concat([df_query, df_query2, df_query3]).reset_index(drop=True)\n",
    "\n",
    "df_query['pdf_page'] = df_query.apply(lambda x: f\"{x.pdf}_{x.page}\", axis=1) \n",
    "df_query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "c5798557-53ba-4aac-801a-aa65a1701814",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  - Recall @1: 0.554\n",
      "  - Recall @3: 0.746\n",
      "  - Recall @5: 0.807\n",
      "  - Recall @10: 0.857\n"
     ]
    }
   ],
   "source": [
    "get_recall_scores(df_query, \"multimodal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5c710f4-ce5a-4308-a0aa-f7231bccd82e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
